#include "core/cell-dyn.hpp"

#include "yocto/code/utils.hpp"
#include "yocto/code/unroll.hpp"
#include "yocto/core/locate.hpp"

#if defined(_OPENMP)
#include <omp.h>
#endif

namespace FastMD
{
	
	static inline int cmp( Cell *lhs, Cell *rhs ) throw()
	{
		return ptrdiff_t(lhs) - ptrdiff_t(rhs);
	}
	
	
	void Cell:: subscribe( Cell *l ) throw()
	{
		assert(links<MAX_LINKS);
		size_t pos = 0;
		if( !core::locate( &l,link,links,pos,cmp) )
		{
			core::insert( &l,link,links,pos);
			size_t &n = (size_t &)links;
			++n;
		}
	}
	
	////////////////////////////////////////////////////////////////////////////
	//
	// Cells
	//
	////////////////////////////////////////////////////////////////////////////
	
	static inline 
	size_t count_cells_for( const Real L, const Real cutout )
	{
		size_t nc = static_cast<size_t>( __floor(L/cutout) );
		if( nc < 1 ) nc = 1;
		return nc;
	}
	
	U3D Cells:: CountFor( const V3D &box, const Real cutout) throw()
	{
		return U3D( count_cells_for(box.x,cutout), count_cells_for(box.y,cutout), count_cells_for(box.z,cutout) );
	}
	
	
	
#define __LINK(INDEX) do { assert( (INDEX) < count );  c.subscribe( &cell[INDEX] );} while(false)
	
	Cells:: Cells( const U3D &dim, const Frame &frame, const size_t num_threads ) :
	size( dim ),
	count( size.items() ),
	side(size.x-1, size.y-1,size.z-1 ),
	ystp( size.x ),
	zstp( size.x * size.y ),
	ibox( frame.ibox ),
	cmem( count ),
	cell( memory::kind<MemoryKind>::acquire_as<Cell>( cmem ) ),
	slices( clamp<size_t>(1,num_threads,MAX_THREADS) ),
	slice()
	{
		
		//----------------------------------------------------------------------
		// prepare slices
		//----------------------------------------------------------------------
		{
			size_t idx = 0;
			size_t num = count;
			for( size_t i=0; i < slices; ++i )
			{
				const size_t ngrp = slices-i;
				const size_t todo = num / ngrp;
				slice[i].ini   = idx;
				slice[i].len   = todo;
				idx += todo;
				slice[i].end   = idx;
				num -= todo;
				std::cerr << "slice[" << i << "]@" << slice[i].ini << "+" << slice[i].len << "<" << slice[i].end << std::endl;
			}
		}
		
		//----------------------------------------------------------------------
		// initialize/create links
		//----------------------------------------------------------------------
		const bool link_x = size.x > 1;
		const bool link_y = size.y > 1;
		const bool link_z = size.z > 1;
		
		const U3D  periodic(1,1,1);
		
		for( size_t k=0; k < size.z; ++k )
		{
			const size_t z_curr   = k * zstp;
			const bool   z_lo     = (k <= 0);
			const bool   z_hi     = (k >= side.z);
			const size_t z_prev   = (z_lo ? side.z : k-1) * zstp;
			const size_t z_next   = (z_hi ? 0      : k+1) * zstp;
			const bool   link_z_prev = link_z && ( (z_lo && periodic.z) || !z_lo );
			const bool   link_z_next = link_z && ( (z_hi && periodic.z) || !z_hi );
			
			for( size_t j=0; j < size.y; ++j )
			{
				const size_t y_curr = j * ystp;
				const bool   y_lo   = (j <= 0);
				const bool   y_hi   = (j >= side.y);
				const size_t y_prev = (y_lo ? side.y : j-1) * ystp;
				const size_t y_next = (y_hi ? 0      : j+1) * ystp;
				const bool   link_y_prev = link_y && ( (y_lo && periodic.y) || !y_lo);
				const bool   link_y_next = link_y && ( (y_hi && periodic.y) || !y_hi);
				
				for( size_t i=0; i < size.x; ++i )
				{
					const bool   x_lo   = (i <= 0);
					const bool   x_hi   = (i >= side.x);
					const size_t x_prev = (x_lo? side.x : i-1);
					const size_t x_next = (x_hi? 0      : i+1);
					const bool   link_x_prev = link_x && ( (x_lo && periodic.x) || !x_lo );
					const bool   link_x_next = link_x && ( (x_hi && periodic.x) || !x_hi );
					
					
#define x_curr i
					//----------------------------------------------------------
					//-- initialize current cell
					//----------------------------------------------------------
					const size_t p = z_curr + y_curr + x_curr;
					const U3D    q(i,j,k);
					Cell        &c = cell[p];
					(size_t&)(c.links) = 0;
					(size_t&)(c.index) = i;
					(U3D   &)(c.id)    = q;
					
					//----------------------------------------------------------
					//-- make 26 tests...
					//----------------------------------------------------------
					//#define __CHECK(X,Y,Z) std::cerr << "(" << X << "," << Y << "," << Z << ")" << std::endl;
#define __CHECK(X,Y,Z)
					
					{
						const bool   lnk = link_x_prev && link_y_prev;
						const size_t ini = x_prev + y_prev;
						__CHECK( -1, -1, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK( -1, -1,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK( -1, -1,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const bool   lnk = link_x_prev;
						const size_t ini = x_prev + y_curr;
						__CHECK( -1,  0, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK( -1,  0,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK( -1,  0,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const bool   lnk = link_x_prev && link_y_next;
						const size_t ini = x_prev + y_next;
						__CHECK( -1,  1, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK( -1,  1,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK( -1,  1,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const bool   lnk = link_y_prev;
						const size_t ini = x_curr + y_prev;
						__CHECK(  0, -1, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK(  0, -1,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK(  0, -1,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const size_t ini = x_curr + y_curr;
						__CHECK(  0,  0, -1 ) if( link_z_prev ) __LINK(ini+z_prev);
						__CHECK(  0,  0,  0 ) // nothing
						__CHECK(  0,  0,  1 ) if( link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const size_t ini = x_curr + y_next;
						const bool   lnk = link_y_next;
						__CHECK(  0,  1, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK(  0,  1,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK(  0,  1,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const size_t ini = x_next + y_prev;
						const bool   lnk = link_x_next && link_y_prev;
						__CHECK(  1, -1, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK(  1, -1,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK(  1, -1,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const size_t ini = x_next + y_curr;
						const bool   lnk = link_x_next;
						__CHECK(  1,  0, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK(  1,  0,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK(  1,  0,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
					
					{
						const size_t ini = x_next + y_next;
						const bool   lnk = link_x_next && link_y_next;
						__CHECK(  1,  1, -1 ) if( lnk && link_z_prev ) __LINK(ini+z_prev);
						__CHECK(  1,  1,  0 ) if( lnk )                __LINK(ini+z_curr);
						__CHECK(  1,  1,  1 ) if( lnk && link_z_next ) __LINK(ini+z_next);
					}
				}
			}
		}
		
		
		
	}
	
	Cells:: ~Cells() throw()
	{
		memory::kind<MemoryKind>::release_as<Cell>( cell, cmem );
	}
	
#define __INDEX(FIELD) min_of<size_t>( max_of<long>( 0, static_cast<long>( size.FIELD * (( r.FIELD * ibox.FIELD ) + __HALF )) ),side.FIELD)
	
	void Cells:: dispatch( CellNode *node, const size_t nodes, const Vector &positions ) throw()
	{
		assert(nodes>=positions.vertices);
#define INIT_CELL(INDEX) cell[ INDEX ].inside = NULL
		YOCTO_LOOP_FUNC_(count,INIT_CELL,0);
		
		const size_t n=positions.vertices;
		for( size_t i=0; i < n; ++i )
		{
			const V3D &r = positions[i];
			const U3D  j( __INDEX(x), __INDEX(y), __INDEX(z) );
			Cell      &c = cell[ j.z * zstp + j.y * ystp + j.x ];
			node->next = c.inside;
			c.inside   = node;
			++node;
		}
		
	}
	
	
	////////////////////////////////////////////////////////////////////////////
	//
	// CellDynamicsBase
	//
	////////////////////////////////////////////////////////////////////////////
	CellDynamicsBase:: ~CellDynamicsBase() throw()
	{
		if(cells) release_cells();
		release_nodes();
	}
	
	CellDynamicsBase:: CellDynamicsBase( const Frame &frame, const Real CutOut, const size_t nt) :
	Dynamics(),
	cutout( CutOut ),
	nmem( frame.particles ),
	node( memory::kind<MemoryKind>::acquire_as<CellNode>( nmem ) ),
	nodes( frame.particles ),
	cells( NULL ),
	num_threads( nt ),
	wksp()
	{
		for( size_t i=0; i < nodes; ++i ) 
			node[i].index=i;
		build_cells(frame);
	}
	
	void CellDynamicsBase:: release_nodes() throw()
	{
		memory::kind<MemoryKind>::release_as<CellNode>( node, nmem );
		(size_t&)nodes = 0;
	}
	
	void CellDynamicsBase:: release_cells() throw()
	{
		assert( cells != NULL );
		cells->~Cells();
		cells = NULL;
	}
	
	void CellDynamicsBase:: dispatch_only( const Vector &positions ) throw()
	{
		assert( cells != NULL );
		cells->dispatch( node,nodes, positions);
	}
	
	void CellDynamicsBase:: build_cells( const Frame &frame )
	{
		const U3D dims = Cells::CountFor( frame.box, cutout);
		
		if( cells )
		{
			//-- update
			if( dims == cells->size )
			{
				(V3D &)(cells->ibox) = frame.ibox;
				return;
			}
			//-- shall rebuild...
			release_cells();
		}
		assert( NULL == cells );
		
		memset( wksp, 0, sizeof(wksp) );              //-- not necessary
		new (wksp) Cells( dims, frame, num_threads ); //-- may throw
		cells = (Cells *) &wksp[0];                   //-- everything went fine
	}
	
	
	
	////////////////////////////////////////////////////////////////////////////
	//
	// CellDynamics
	//
	////////////////////////////////////////////////////////////////////////////
	CellDynamics:: ~CellDynamics() throw()
	{
		
	}
	
	
	CellDynamics:: CellDynamics( const Frame &frame, const Real CutOut ) :
	CellDynamicsBase(frame,CutOut,1)
	{
	}
	
	void CellDynamics:: eval( Frame &frame, Potential &potential, bool isothermal ) throw()
	{
		prolog(frame);
		work(0,frame,potential);
		epilog(frame,isothermal);
	}
	
	void CellDynamicsBase:: prolog( Frame &frame ) throw()
	{
		//======================================================================
		//
		// dispatch all positions
		//
		//======================================================================
		cells->dispatch( node,nodes,frame[ Frame::POSITIONS ]);
		
		//======================================================================
		//
		// initialize forces/accelerations
		//
		//======================================================================
		frame[ Frame::ACCELERATIONS ].ldz();
	}
	
	void CellDynamicsBase:: work( size_t rank, Frame &frame, Potential &potential ) throw()
	{
		Cells        &C        = *cells;
		const Vector &position =  frame[ Frame::POSITIONS     ];
		Vector       &force    =  frame[ Frame::ACCELERATIONS ];
		PBC_Proc      pbc_f    =  frame.pbc_f;
		const V3D    &box      =  frame.box;
		const V3D    &ibox     =  frame.ibox;
		
		assert( rank <  cells->slices);
		const size_t c_ini = C.slice[rank].ini; assert(c_ini<cells->count);
		const size_t c_end = C.slice[rank].end; assert(c_end<=cells->count);
		
		for( size_t ic=c_ini; ic < c_end; ++ic )
		{
			Cell         &c    = C[ic];
			Cell        **link = c.link;
			const size_t  nl   = c.links;
			
			//--------------------------------------------------------------
			// initialize local potential
			//--------------------------------------------------------------
			Real V = 0;
			V3D  f_ij;
			
			//--------------------------------------------------------------
			// outer loop on inside particles
			//--------------------------------------------------------------
			CellNode *in = c.inside;
			while( in )
			{
				const size_t i   = in->index;
				const V3D   &r_i = position[i];
				V3D         &f_i = force[i];
				
				//----------------------------------------------------------
				// First inner loop: inside the same cell
				//----------------------------------------------------------
				CellNode *jn = in->next;
				while( jn )
				{
					const size_t j   = jn->index; assert(i!=j);
					const V3D   &r_j = position[j];
					V3D          r_ij(r_i,r_j);
					V3D         &f_j = force[j];
					
					pbc_f(r_ij,box,ibox);
					const Real V_ij = potential(f_ij,r_ij);
					V += (V_ij + V_ij);
					f_i.sub( f_ij );
					f_j.add( f_ij );
					jn = jn->next;
				}
				
				//----------------------------------------------------------
				// Second inner loop: over the neighbors
				//----------------------------------------------------------
				for( size_t il=0; il < nl; ++il )
				{
					assert( link[il] != NULL );
					jn = link[il]->inside;					
					while( jn )
					{
						//--------------------------------------------------
						// update V and only f_i
						//--------------------------------------------------
						const size_t j = jn->index;
						const V3D &r_j = position[j];
						V3D        r_ij(r_i,r_j);
						
						pbc_f(r_ij,box,ibox);
						V += potential(f_ij,r_ij);
						f_i.sub( f_ij );
						jn = jn->next;
					}
				}
				
				in = in->next;
			}
			c.V = V;
		}	
		
		
	}
	
	void CellDynamicsBase:: epilog( Frame &frame, bool isothermal ) throw()
	{
		//======================================================================
		//
		// collect data
		//
		//======================================================================
		const Cells &C  = *cells;
		const size_t nc = C.count;
		frame.V = 0;
		for( size_t ic=0; ic < nc; ++ic )
		{
			frame.V += C[ic].V;
		}
		frame.V *= __HALF;
		
		if( isothermal )
			frame.thermal_friction();
		frame.forces_to_accelerations();
		
	}
	
	
		
	
	
}
